package com.steak.agent;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

/**
 * @author 刘牌
 * @Title:
 * @date 2022/3/1722:48
 */
public class HelloAgent {
    public static void premain(String ops , Instrumentation instrumentation) {
        System.out.println("premain start , params is "+ops);
        DefineTransformer defineTransformer = new DefineTransformer();
        instrumentation.addTransformer(defineTransformer);
    }

    static final Set<String> classNameSet = new HashSet<>();

    static {
        classNameSet.add("com.steak.agent.AgentTest");
    }

    static class DefineTransformer implements ClassFileTransformer{
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
           try {
               String currentClassName = className.replaceAll("/", ".");
               if (!classNameSet.contains(currentClassName)) return null;
               CtClass ctClass = ClassPool.getDefault().get(currentClassName);
               CtBehavior[] declaredBehaviors = ctClass.getDeclaredBehaviors();
               for (CtBehavior declaredBehavior : declaredBehaviors) {
                    enhanceMethod(declaredBehavior);
               }
               return ctClass.toBytecode();
           }catch (Exception e){
               e.printStackTrace();
           }
           return null;
        }
    }

    private static void enhanceMethod(CtBehavior method) throws CannotCompileException {
        if (method.isEmpty()) return;
        String methodName = method.getName();
        if ("main".equalsIgnoreCase(methodName)) return;
        final StringBuilder stringBuilder = new StringBuilder();
        stringBuilder.append("{")
                .append("long start = System.nanoTime();\n") //前置增强: 打入时间戳
                .append("$_ = $proceed($$);\n")              //调用原有代码，类似于method();($$)表示所有的参数
                .append("System.out.print(\"method:[")
                .append(methodName).append("]\");").append("\n")
                .append("System.out.println(\" cost:[\" +(System.nanoTime() - start)+ \"ns]\");") // 后置增强，计算输出方法执行耗时
                .append("}");
        ExprEditor exprEditor = new ExprEditor() {
            @Override
            public void edit(MethodCall m) throws CannotCompileException {
                m.replace(stringBuilder.toString());
            }
        };
        method.instrument(exprEditor);
    }
}
